from py2neo import Graph, NodeMatcher, RelationshipMatcher
from queue import Queue
import argparse

DEFAULT_HOST = "211.71.15.39"
DEFAULT_PORT = "9474"
DEFAULT_USERNAME = "python"
DEFAULT_PASSWORD = "python"
DEFAULT_PKG_NAME = "Package"
DEFAULT_VUL_NAME = "Vulnerability"
DEFAULT_PKG_PKG_RELATION = "depends_on_pkg"
DEFAULT_PKG_VUL_RELATION = "affect_pkg"

pkg_queue = Queue()
path = []
vul_set = set()

class DataBase:
    def __init__(self, url, user, pwd):
        self.graph = Graph(url, password=pwd, name=user)
        self.nodes = NodeMatcher(self.graph)
        self.relations = RelationshipMatcher(self.graph)

    # 通过id查找包
    def match_pkg_by_id(self, pkg_name, id):
        res = self.nodes.match(pkg_name, ID=id)
        return res.first()

    # 查询结点出发的特定关系
    def query_relations_by_node1(self, node, relation=None):
        res = list(self.relations.match([node], r_type=relation))
        return res

    # 查询到达结点的特定关系
    def query_relations_by_node2(self, node, relation=None):
        res = list(self.relations.match([None, node], r_type=relation))
        return res

class RTree:
    def __init__(self, id):
        self.id = id
        self.children = []
        self.vuls = []
        self.vuls_all = []

    def get_id(self):
        return self.id

    def add_child(self, node):
        self.children.append(node)

    def add_vuls(self, relations):
        for r in relations:
            vul = r.start_node
            if len(vul['CVE_ID']) > 0:
                self.vuls.append(vul['CVE_ID'])
            else:
                self.vuls.append(vul['ID'])

    # dfs遍历树并按先序打印路径
    def print_path(self):
        if len(self.vuls) == 0:
            path.append(self.id)
        else:
            str = self.id + '('
            for index in range(len(self.vuls)):
                if index < len(self.vuls) - 1:
                    str += self.vuls[index] + ','
                else:
                    str += self.vuls[index] + ')'
            path.append(str)
        if len(self.children) == 0:
            for index in range(len(path)):
                if index < len(path) - 1:
                    print(path[index], end = "->")
                else:
                    print(path[index])
            print()
            path.pop()
            return

        for child in self.children:
            child.print_path()

        path.pop()
        return

    def print_vuls(self):
        node_queue = Queue()
        node_queue.put(self)
        while not node_queue.empty():
            cur = node_queue.get()
            if len(cur.vuls):
                for vul in cur.vuls:
                    vul_set.add(vul)
            if len(cur.children):
                for child in cur.children:
                    node_queue.put(child)
        vul_list = list(vul_set)
        for vul in vul_list:
            print(vul)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()  # 创建一个解析对象
    # 向该对象中添加要关注的命令行参数和选项
    parser.add_argument('--host', default=DEFAULT_HOST)
    parser.add_argument('--port', default=DEFAULT_PORT)
    parser.add_argument('--username', default=DEFAULT_USERNAME)
    parser.add_argument('--password', default=DEFAULT_PASSWORD)
    parser.add_argument('--pkg', default=DEFAULT_PKG_NAME)
    parser.add_argument('--vul', default=DEFAULT_VUL_NAME)
    parser.add_argument('--pkg_pkg', default=DEFAULT_PKG_PKG_RELATION)
    parser.add_argument('--pkg_vul', default=DEFAULT_PKG_VUL_RELATION)
    parser.add_argument('--rela_type', default=2)
    args = parser.parse_args()  # 进行解析

    url = 'http://' + args.host + ':' + args.port
    username = args.username
    password = args.password
    db = DataBase(url, username, password)

    pkg_node = args.pkg
    vul_node = args.vul
    pkg_pkg = args.pkg_pkg
    pkg_vul = args.pkg_vul
    rela_type = args.rela_type

    print("please input the pkg-id which you are searching for:")
    input_pkg = input()
    root = RTree(input_pkg)
    pkg_queue.put(root)
    while not pkg_queue.empty():
        cur = pkg_queue.get()
        node1 = db.match_pkg_by_id(pkg_node, cur.get_id())
        if node1 is None:
            print("%s is not in database" % cur.get_id())
            break
        if rela_type == 2:  # vul -> pkg
            vuls = db.query_relations_by_node2(node1, pkg_vul)
        elif rela_type == 1:    # pkg -> vul
            vuls = db.query_relations_by_node1(node1, pkg_vul)
        else:
            print("rela_type is illeagal")
            break
        if len(vuls) != 0:
            cur.add_vuls(vuls)
        rela_list = db.query_relations_by_node1(node1, pkg_pkg)
        for relation in rela_list:
            node2 = relation.end_node
            nxt = RTree(node2['ID'])
            cur.add_child(nxt)
            pkg_queue.put(nxt)

    root.print_vuls()